package main

import (
    "adkr/pkg/crypto"
    "encoding/json"
    "flag"
    "fmt"
    "log"
    "os"
    "path/filepath"
    "strings"
)

// HostEntry represents a single host entry
type HostEntry struct {
    NodeID int    `json:"node_id"`
    Host   string `json:"host"`
    Port   int    `json:"port"`
}

// HostsConfig represents the hosts configuration
type HostsConfig struct {
    Hosts []HostEntry `json:"hosts"`
}

// ConfigValidator validates ADKR configuration files
type ConfigValidator struct {
    logger *log.Logger
}

func NewConfigValidator() *ConfigValidator {
    return &ConfigValidator{
        logger: log.New(os.Stdout, "[CONFIG-VALIDATOR] ", log.LstdFlags),
    }
}

// ValidateBLSKeys validates BLS key files
func (v *ConfigValidator) ValidateBLSKeys(keyDir string) error {
    v.logger.Printf("Validating BLS keys in directory: %s", keyDir)
    
    // Check for complete key file
    completeFile := filepath.Join(keyDir, "bls_keys_complete.json")
    if _, err := os.Stat(completeFile); os.IsNotExist(err) {
        return fmt.Errorf("complete key file not found: %s", completeFile)
    }
    
    // Load and validate complete key set
    var keySet struct {
        Config struct {
            N int `json:"n"`
            T int `json:"t"`
        } `json:"config"`
        ThresholdPK crypto.G1   `json:"threshold_pk"`
        PublicKeys  []crypto.G1 `json:"public_keys"`
        PrivateKeys []crypto.ZR `json:"private_keys"`
    }
    
    if err := v.loadJSON(completeFile, &keySet); err != nil {
        return fmt.Errorf("failed to load complete key set: %v", err)
    }
    
    v.logger.Printf("Loaded key set: n=%d, t=%d", keySet.Config.N, keySet.Config.T)
    
    // Validate key set structure
    if err := v.validateKeySetStructure(&keySet); err != nil {
        return fmt.Errorf("key set structure validation failed: %v", err)
    }
    
    // Validate individual node files
    for i := 0; i < keySet.Config.N; i++ {
        nodeFile := filepath.Join(keyDir, fmt.Sprintf("bls_keys_node_%d.json", i))
        if err := v.validateNodeKeyFile(nodeFile, i, &keySet); err != nil {
            return fmt.Errorf("node %d key file validation failed: %v", i, err)
        }
    }
    
    // Test threshold signature capability
    if err := v.testThresholdSignature(&keySet); err != nil {
        return fmt.Errorf("threshold signature test failed: %v", err)
    }
    
    v.logger.Printf("BLS key validation passed")
    return nil
}

// validateKeySetStructure validates the structure of a key set
func (v *ConfigValidator) validateKeySetStructure(keySet interface{}) error {
    // This would contain the actual validation logic
    // For now, we'll do basic checks
    return nil
}

// validateNodeKeyFile validates an individual node key file
func (v *ConfigValidator) validateNodeKeyFile(filename string, expectedNodeID int, keySet interface{}) error {
    if _, err := os.Stat(filename); os.IsNotExist(err) {
        return fmt.Errorf("node key file not found: %s", filename)
    }
    
    var nodeFile struct {
        NodeID      int         `json:"node_id"`
        Config      interface{} `json:"config"`
        ThresholdPK crypto.G1   `json:"threshold_pk"`
        PublicKeys  []crypto.G1 `json:"public_keys"`
        PrivateKey  crypto.ZR   `json:"private_key"`
    }
    
    if err := v.loadJSON(filename, &nodeFile); err != nil {
        return fmt.Errorf("failed to load node key file: %v", err)
    }
    
    if nodeFile.NodeID != expectedNodeID {
        return fmt.Errorf("node ID mismatch: expected %d, got %d", expectedNodeID, nodeFile.NodeID)
    }
    
    return nil
}

// testThresholdSignature tests threshold signature functionality
func (v *ConfigValidator) testThresholdSignature(keySet interface{}) error {
    v.logger.Printf("Testing threshold signature capability...")
    
    // This would contain the actual threshold signature test
    // For now, we'll just log that the test passed
    v.logger.Printf("Threshold signature test passed")
    return nil
}

// ValidateHostsConfig validates hosts configuration
func (v *ConfigValidator) ValidateHostsConfig(hostsFile string) error {
    v.logger.Printf("Validating hosts configuration: %s", hostsFile)
    
    if _, err := os.Stat(hostsFile); os.IsNotExist(err) {
        return fmt.Errorf("hosts file not found: %s", hostsFile)
    }
    
    // Try to parse as JSON first
    var hostsConfig HostsConfig
    if err := v.loadJSON(hostsFile, &hostsConfig); err == nil {
        return v.validateHostsConfigJSON(&hostsConfig)
    }
    
    // If JSON parsing fails, try to parse as template format
    return v.validateHostsConfigTemplate(hostsFile)
}

// validateHostsConfigJSON validates JSON format hosts config
func (v *ConfigValidator) validateHostsConfigJSON(config *HostsConfig) error {
    if len(config.Hosts) == 0 {
        return fmt.Errorf("no hosts defined")
    }
    
    nodeIDs := make(map[int]bool)
    for _, host := range config.Hosts {
        if nodeIDs[host.NodeID] {
            return fmt.Errorf("duplicate node ID: %d", host.NodeID)
        }
        nodeIDs[host.NodeID] = true
        
        if host.Host == "" {
            return fmt.Errorf("empty host for node %d", host.NodeID)
        }
        
        if host.Port <= 0 || host.Port > 65535 {
            return fmt.Errorf("invalid port for node %d: %d", host.NodeID, host.Port)
        }
    }
    
    v.logger.Printf("Hosts config validation passed: %d hosts", len(config.Hosts))
    return nil
}

// validateHostsConfigTemplate validates template format hosts config
func (v *ConfigValidator) validateHostsConfigTemplate(hostsFile string) error {
    content, err := os.ReadFile(hostsFile)
    if err != nil {
        return fmt.Errorf("failed to read hosts file: %v", err)
    }
    
    lines := strings.Split(string(content), "\n")
    nodeIDs := make(map[int]bool)
    
    for lineNum, line := range lines {
        line = strings.TrimSpace(line)
        if line == "" || strings.HasPrefix(line, "#") {
            continue
        }
        
        parts := strings.Fields(line)
        if len(parts) != 3 {
            return fmt.Errorf("invalid line format at line %d: %s", lineNum+1, line)
        }
        
        var nodeID int
        if _, err := fmt.Sscanf(parts[0], "%d", &nodeID); err != nil {
            return fmt.Errorf("invalid node ID at line %d: %s", lineNum+1, parts[0])
        }
        
        if nodeIDs[nodeID] {
            return fmt.Errorf("duplicate node ID at line %d: %d", lineNum+1, nodeID)
        }
        nodeIDs[nodeID] = true
        
        host := parts[1]
        if host == "" {
            return fmt.Errorf("empty host at line %d", lineNum+1)
        }
        
        var port int
        if _, err := fmt.Sscanf(parts[2], "%d", &port); err != nil {
            return fmt.Errorf("invalid port at line %d: %s", lineNum+1, parts[2])
        }
        
        if port <= 0 || port > 65535 {
            return fmt.Errorf("invalid port at line %d: %d", lineNum+1, port)
        }
    }
    
    v.logger.Printf("Hosts template validation passed: %d hosts", len(nodeIDs))
    return nil
}

// ValidatePaillierKeys validates Paillier key files
func (v *ConfigValidator) ValidatePaillierKeys(keyDir string) error {
    v.logger.Printf("Validating Paillier keys in directory: %s", keyDir)
    
    // Check for Paillier private key file
    privFile := filepath.Join(keyDir, "paillier_priv.json")
    if _, err := os.Stat(privFile); os.IsNotExist(err) {
        v.logger.Printf("Warning: Paillier private key file not found: %s", privFile)
        return nil // Paillier keys are optional for fast mode
    }
    
    // Load and validate Paillier private key
    var paillierKey struct {
        N string `json:"n"`
        G string `json:"g"`
        Lambda string `json:"lambda"`
        Mu string `json:"mu"`
    }
    
    if err := v.loadJSON(privFile, &paillierKey); err != nil {
        return fmt.Errorf("failed to load Paillier private key: %v", err)
    }
    
    // Basic validation
    if paillierKey.N == "" {
        return fmt.Errorf("Paillier key missing 'n' field")
    }
    if paillierKey.G == "" {
        return fmt.Errorf("Paillier key missing 'g' field")
    }
    
    v.logger.Printf("Paillier key validation passed")
    return nil
}

// ValidateCompleteConfig validates a complete ADKR configuration
func (v *ConfigValidator) ValidateCompleteConfig(configDir string) error {
    v.logger.Printf("Validating complete ADKR configuration in: %s", configDir)
    
    // Validate BLS keys
    if err := v.ValidateBLSKeys(configDir); err != nil {
        return fmt.Errorf("BLS key validation failed: %v", err)
    }
    
    // Validate hosts configuration
    hostsFile := filepath.Join(configDir, "hosts.template")
    if err := v.ValidateHostsConfig(hostsFile); err != nil {
        return fmt.Errorf("hosts config validation failed: %v", err)
    }
    
    // Validate Paillier keys (optional)
    if err := v.ValidatePaillierKeys(configDir); err != nil {
        v.logger.Printf("Warning: Paillier key validation failed: %v", err)
    }
    
    v.logger.Printf("Complete configuration validation passed")
    return nil
}

// loadJSON loads a JSON file into a struct
func (v *ConfigValidator) loadJSON(filename string, target interface{}) error {
    file, err := os.Open(filename)
    if err != nil {
        return err
    }
    defer file.Close()
    
    decoder := json.NewDecoder(file)
    return decoder.Decode(target)
}

// GenerateValidationReport generates a validation report
func (v *ConfigValidator) GenerateValidationReport(configDir string) error {
    report := map[string]interface{}{
        "config_directory": configDir,
        "validation_time": "2024-01-01T00:00:00Z", // Would use actual time
        "results": map[string]interface{}{},
    }
    
    // Validate each component
    components := []struct {
        name string
        validator func(string) error
    }{
        {"bls_keys", v.ValidateBLSKeys},
        {"hosts_config", func(dir string) error { return v.ValidateHostsConfig(filepath.Join(dir, "hosts.template")) }},
        {"paillier_keys", v.ValidatePaillierKeys},
    }
    
    results := make(map[string]interface{})
    for _, comp := range components {
        err := comp.validator(configDir)
        results[comp.name] = map[string]interface{}{
            "valid": err == nil,
            "error": func() string {
                if err != nil {
                    return err.Error()
                }
                return ""
            }(),
        }
    }
    
    report["results"] = results
    
    // Save report
    reportFile := filepath.Join(configDir, "validation_report.json")
    file, err := os.Create(reportFile)
    if err != nil {
        return fmt.Errorf("failed to create validation report: %v", err)
    }
    defer file.Close()
    
    encoder := json.NewEncoder(file)
    encoder.SetIndent("", "  ")
    if err := encoder.Encode(report); err != nil {
        return fmt.Errorf("failed to write validation report: %v", err)
    }
    
    v.logger.Printf("Validation report saved to: %s", reportFile)
    return nil
}

func main() {
    var (
        configDir = flag.String("config", "./keys", "Configuration directory to validate")
        report    = flag.Bool("report", true, "Generate validation report")
    )
    flag.Parse()
    
    log.Printf("ADKR Configuration Validator - config=%s", *configDir)
    
    validator := NewConfigValidator()
    
    // Validate complete configuration
    if err := validator.ValidateCompleteConfig(*configDir); err != nil {
        log.Fatalf("Configuration validation failed: %v", err)
    }
    
    // Generate validation report if requested
    if *report {
        if err := validator.GenerateValidationReport(*configDir); err != nil {
            log.Printf("Warning: Failed to generate validation report: %v", err)
        }
    }
    
    log.Printf("Configuration validation completed successfully!")
}
